from langchain_community.vectorstores import Chroma
from typing import Dict, Any
from env import (
    model_url,
    model_name,
    collection_name,
    collection_dir,
    mongodb_uri,
    db_name,
)


# 初始化 Chroma 集合
def init_chroma(embedding_model):
    return Chroma(
        collection_name=collection_name,
        embedding_function=embedding_model,
        persist_directory=collection_dir,
    )


def extract_list(val: str) -> list[str]:
    """
    从val中获取metadata，并将其转换成数组
    """

    start_tag = "metadata-start\n"
    end_tag = "metadata-end\n"

    start_index = val.find(start_tag)
    if start_index == -1:
        return []

    # Move the index past the start tag
    start_index += len(start_tag)

    end_index = val.find(end_tag, start_index)
    if end_index == -1:
        return []

    if start_index == end_index:
        return []

    # Extract the substring between the start and end tags
    meta_data = val[start_index:end_index].strip()

    return meta_data.split("\n")


def build_list_to_json(items: list[str]) -> Dict[str, Any]:
    result = {}
    for item in items:
        # 去除首尾空格
        item = item.strip()
        # 通过冒号分割键值对
        key, value = item.split(":", 1)
        # 去除键和值的首尾空格
        key = key.strip()
        value = value.strip()
        # 尝试将值转换为合适的类型，例如整数或浮点数
        try:
            value = int(value)  # 尝试转换为整数
        except ValueError:
            try:
                value = float(value)  # 尝试转换为浮点数
            except ValueError:
                pass  # 如果无法转换，则保持为字符串
        result[key] = value
    return result
